from sqlalchemy.orm import Session
from schema import AI_Model_schema
from model import AI_Model_model



def get_AI_Model(db: Session, ID: int):
    return db.query(AI_Model_model.Ai_model).filter(AI_Model_model.Ai_model.ID == ID).first()


def get_all_AI_Model(db: Session):
    return db.query(AI_Model_model.Ai_model).all()


def delete_AI_Model(db: Session, ID):
    to_delete = db.query(AI_Model_model.Ai_model).filter(AI_Model_model.Ai_model.ID == ID)
    deleted = to_delete.delete()
    if deleted == 0:
        return None
    else:
        db.commit()
        return True


def create_AI_Model(db: Session, AI_Model: AI_Model_schema.Ai_model):
    db_AI_Model = AI_Model_model.Ai_model(
        ModelName=AI_Model.ModelName,
        ModelFunction=AI_Model.ModelFunction,
        InputType=AI_Model.InputType,
        OutputType=AI_Model.OutputType,
        ModelAPIAddress=AI_Model.ModelAPIAddress,
    )
    db.add(db_AI_Model)
    db.commit()
    db.refresh(db_AI_Model)
    return db_AI_Model


def update_AI_Model(db: Session, ID, field_name, new_value):
    db_AI_Model = db.query(AI_Model_model.Ai_model).filter(AI_Model_model.Ai_model.ID == ID).first()
    setattr(db_AI_Model, field_name, new_value)
    db.commit()
    db.refresh(db_AI_Model)
    return db_AI_Model
